#
# Copyright (C) [2024] Xingyun Integrated Circuit, Inc.
#
# GreenCode was a private technology asset of Xingyun Integrated Circuit， Inc （Confidential）
# Author: Shawn.Tan
# Date: 2025.10.28
#
# History: Initial Version 2025.10.28
#
#
# Basic test for speculative BNT sequence

from gpgpu.EnvGPGPU import EnvGPGPU
from gpgpu.GenThreadGPGPU import GenThreadGPGPU
from base.Sequence import Sequence
from base.BntSequence import BntSequence
from base.ChoicesModifier import *


class MyBntSequence(BntSequence):
    def __init__(self, gen_thread, instr_num=10):
        super(MyBntSequence, self).__init__(gen_thread, instr_num)

    def BntSequence_IntegerArithmetic(self):
        self.notice("Speculative Generate Bnt Sequence for Arithmetic instructions")

        if self.getGlobalState("AppRegisterWidth") == 32:
            math_instr_dict = {
                "ADD##CALM": 10,
                "ADDI##CALM": 10,
                "SUB##CALM": 10,
                "MUL##CALM": 10,
                "MULH##CALM": 10,
                "MULHSU##CALM": 10,
                "MULHU##CALM": 10,
                "DIV##CALM": 10,
                "DIVU##CALM": 10,
            }
        else:
            math_instr_dict = {
                "ADD##CALM": 10,
                "ADDW##CALM": 10,
                "ADDI##CALM": 10,
                "ADDIW##CALM": 10,
                "SUB##CALM": 10,
                "SUBW##CALM": 10,
                "MUL##CALM": 10,
                "MULH##CALM": 10,
                "MULHSU##CALM": 10,
                "MULHU##CALM": 10,
                "MULW##CALM": 10,
                "DIV##CALM": 10,
                "DIVU##CALM": 10,
                "DIVUW##CALM": 10,
                "DIVW##CALM": 10,
            }

        for i in range(self.instr_num):
            instr_chosen = self.pickWeighted(math_instr_dict)
            self.genInstruction(instr_chosen)

        # set choices modifier shared by both taken and not-taken path
        choice_mod = ChoicesModifier(self.genThread)
        choice_mod.modifyOperandChoices(
            "GPRs", {"x1": 10000, "x2": 0}
        )  # X2 should be stack pointer
        mod_id = choice_mod.commitSet()

        self.setBntHook(
            **{"Func": "BntSequence_IntegerArithmetic"}
        )  # set the bnt sequence for the following branch instructions
        self.genInstruction(
            "JALR##CALM",
            {
                "rs1": 1,
                "BRarget": "0xc0000000-0xc80000000",
                "SpeculativeBnt": True,
            },
        )
        self.genMetaInstruction("ADDI##CALM", {"simm12": 0x200})  # instructions on taken path
        self.revertBntHook()  # revert the bnt hook to the last setting

        choice_mod.revert(mod_id)

    def BntSequence_L(self):
        self.notice("Speculative Generate Bnt Sequence for L")

        if self.getGlobalState("AppRegisterWidth") == 32:
            ls_dict = {
                "LR.W##CALM": 10,
                "LUI##CALM": 10,
                "LW##CALM": 10,
                "LB##CALM": 10,
                "LBU##CALM": 10,
                "LH##CALM": 10,
                "LHU##CALM": 10,
                "SC.W##CALM": 10,
                "SW##CALM": 10,
                "SH##CALM": 10,
                "SB##CALM": 10,
            }
        else:
            ls_dict = {
                "LR.D##CALM": 10,
                "LR.W##CALM": 10,
                "LUI##CALM": 10,
                "LW##CALM": 10,
                "LWU##CALM": 10,
                "LB##CALM": 10,
                "LBU##CALM": 10,
                "LH##CALM": 10,
                "LHU##CALM": 10,
                "LD##CALM": 10,
                "SC.D##CALM": 10,
                "SC.W##CALM": 10,
                "SW##CALM": 10,
                "SD##CALM": 10,
                "SH##CALM": 10,
                "SB##CALM": 10,
            }

        for i in range(self.instr_num):
            instr_chosen = self.pickWeighted(ls_dict)
            self.genInstruction(instr_chosen, {"LSTarget": "0xd0000000-0xe0000000"})
        cur_pc = self.getPEstate("PC")
        self.genInstruction("JAL##CALM", {"BRTarget": cur_pc})  # branch to itself


class MainSequence(Sequence):
    def generate(self, **kargs):

        # set choices modifier shared by both taken and not-taken paths
        choice_mod = ChoicesModifier(self.genThread)
        choice_mod.modifyOperandChoices(
            "GPRs", {"x0": 10000, "x2": 0}
        )  # X2 should be stack pointer
        mod_id = choice_mod.commitSet()

        id_arithmetic = self.setBntHook(
            Seq="MyBntSequence", Func="BntSequence_IntegerArithmetic"
        )  # set the bnt function for the following branch instructions
        self.genInstruction(
            "BGE##CALM", {"SpeculativeBnt": True}
        )  # some conditional branch instruction
        self.genMetaInstruction("ADDI##CALM", {"simm12": 0x100})  # instructions on taken path

        # revert the modifier
        choice_mod.revert(mod_id)

        id_l = self.setBntHook(
            Func="BntSequence_L"
        )  # set the bnt function for the following branch instructions
        self.genInstruction("JALR##CALM", {"SpeculativeBnt": True})
        self.genMetaInstruction("ADDI##CALM")  # instructions on the taken path

        self.revertBntHook(id_l)  # revert to the last setting: L-->Arithmetic


# Points to the MainSequence defined in this file
MainSequenceClass = MainSequence

# Using the GenThreadGPGPU by default, can be overwritten
GenThreadClass = GenThreadGPGPU

# Using EnvGPGPU by default
EnvClass = EnvGPGPU
